import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
  def __init__(self, d_model, dropout=0.1, max_len=5000):
    # d_model: embedding维度d
    super(PositionalEncoding, self).__init__()
    self.dropout = nn.Dropout(p=dropout)

    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # max_len, 1
    # [10000^(2i/d_model)]^(-1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)  # (max_len, 1) * (d_model/2,) -> (max_len, d_model/2)
    pe[:, 1::2] = torch.cos(position * div_term)
    pe = pe.unsqueeze(0)  # (1, max_len, d_model)
    self.register_buffer('pe', pe)

  def forward(self, x):
    # (N, seq_len, d_model)
    batch_size, seq_len, _ = x.size()
    pe = self.pe[:, :seq_len]
    pe = pe.expand(batch_size, -1, -1)
    x = x + pe
    return self.dropout(x)

class TransformerEncoderLayer(nn.Module):
  # ... (TransformerEncoderLayer definition)
    pass

class TransformerEncoder(nn.Module):
  def __init__(self, encoder_layer, num_layers, norm=None):
    super(TransformerEncoder, self).__init__()
    self.layers = _get_clones(encoder_layer, num_layers)
    self.num_layers = num_layers
    self.norm = norm
    self.pos_encoder = PositionalEncoding(d_model)

  def forward(self, src, src_mask=None, src_key_padding_mask=None):
    src = self.pos_encoder(src)  # 添加位置编码
    output = src
    for mod in self.layers:
      output = mod(output, src_mask=src_mask,
                   src_key_padding_mask=src_key_padding_mask)

    if self.norm is not None:
      output = self.norm(output)

    return output

class TransformerDecoderLayer(nn.Module):
  # ... (TransformerDecoderLayer definition)

class TransformerDecoder(nn.Module):
  def __init__(self, decoder_layer, num_layers, norm=None):
    super(TransformerDecoder, self).__init__()
    self.layers = _get_clones(decoder_layer, num_layers)
    self.num_layers = num_layers
    self.norm = norm
    self.pos_encoder = PositionalEncoding(d_model)

  def forward(self, tgt, memory, tgt_mask=None,
              memory_mask=None, tgt_key_padding_mask=None,
              memory_key_padding_mask=None):
    tgt = self.pos_encoder(tgt)  # 添加位置编码
    output = tgt
    for mod in self.layers:
      output = mod(output, memory, tgt_mask=tgt_mask,
                   memory_mask=memory_mask,
                   tgt_key_padding_mask=tgt_key_padding_mask,
                   memory_key_padding_mask=memory_key_padding_mask)

    if self.norm is not None:
      output = self.norm(output)

    return output

# ... (Rest of the code)